import torch
import matplotlib.pyplot as plt
from PIL import Image
from torchvision import transforms


#  绘制常用激活函数的曲线
def plot_activation_functions():
    """
    绘制常用激活函数的曲线（一维输入）
    """
    # 1. 生成一维输入（从-5到5，包含正负值）
    x = torch.linspace(-5, 5, 100)  # 一维张量，100个点

    # 2. 计算各激活函数的输出
    y_relu = torch.relu(x)
    y_sigmoid = torch.sigmoid(x)
    y_tanh = torch.tanh(x)

    # 3. 转换为numpy（用于matplotlib绘图）
    x_np = x.numpy()
    y_relu_np = y_relu.numpy()
    y_sigmoid_np = y_sigmoid.numpy()
    y_tanh_np = y_tanh.numpy()

    # 4. 绘制图像
    plt.figure(figsize=(15, 4))

    plt.subplot(1, 3, 1)
    plt.plot(x_np, y_relu_np)
    plt.title("ReLU Activation")
    plt.grid()

    plt.subplot(1, 3, 2)
    plt.plot(x_np, y_sigmoid_np)
    plt.title("Sigmoid Activation")
    plt.grid()

    plt.subplot(1, 3, 3)  # 创建1行3列的子图网格，并选中第3个子图
    plt.plot(x_np, y_tanh_np)  # 绘制Tanh函数曲线：x_np为输入值，y_tanh_np为对应的Tanh输出值
    plt.title("Tanh Activation")  # 设置子图标题为"Tanh Activation"
    plt.grid()  # 显示网格线（便于观察数值比例）

    plt.tight_layout()
    plt.show()


# 激活函数应用于图像
def visualize_activations(image_path):
    """
    用实际图像展示不同激活函数的效果
    参数:
        image_path: 输入图像的路径
    """
    # 1. 读取图像并转为灰度图
    img = Image.open(image_path).convert('L')

    # 2. 图像预处理（转为Tensor并归一化）
    transform = transforms.Compose([
        transforms.ToTensor(),  # 自动归一化到[0,1]
    ])
    img_tensor = transform(img).unsqueeze(0)  # 增加batch维度 [1,1,H,W]

    # 3. 应用不同激活函数
    with torch.no_grad():
        relu_out = torch.relu(img_tensor)  # ReLU处理
        sigmoid_out = torch.sigmoid(img_tensor)  # Sigmoid处理
        tanh_out = torch.tanh(img_tensor * 2 - 1)  # Tanh需要调整输入范围到[-1,1]

    # 4. 转换为可显示的图像格式
    def tensor_to_image(tensor):
        return transforms.ToPILImage()(tensor.squeeze())

    original_img = tensor_to_image(img_tensor)
    relu_img = tensor_to_image(relu_out)
    sigmoid_img = tensor_to_image(sigmoid_out)
    tanh_img = tensor_to_image(tanh_out + 1)  # Tanh输出[-1,1]转回[0,1]

    # 5. 可视化对比
    plt.figure(figsize=(16, 4))

    # 原始图像
    plt.subplot(1, 4, 1)
    plt.imshow(original_img, cmap='gray')
    plt.title("Original Image")
    plt.axis('off')

    # ReLU处理结果
    plt.subplot(1, 4, 2)
    plt.imshow(relu_img, cmap='gray')
    plt.title("ReLU Activation")
    plt.axis('off')

    # Sigmoid处理结果
    plt.subplot(1, 4, 3)
    plt.imshow(sigmoid_img, cmap='gray')
    plt.title("Sigmoid Activation")
    plt.axis('off')

    # Tanh处理结果
    plt.subplot(1, 4, 4)
    plt.imshow(tanh_img, cmap='gray')
    plt.title("Tanh Activation")
    plt.axis('off')

    plt.tight_layout()
    plt.show()


if __name__ == '__main__':
    # plot_activation_functions()
    visualize_activations('../images/dog.png')